import os

import torch
from tqdm import tqdm
from eval_metrics.metrics import cal_acc_precision_recall_f1

class Trainer:
    def __init__(self, model, optimizer, criterion, scheduler=None, device=None):
        """
        初始化Trainer

        :param model: 需要训练的模型
        :param optimizer: 优化器
        :param criterion: 损失函数
        :param device: 训练使用的设备（默认为cuda或cpu）
        """
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.scheduler = scheduler
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        self.best_metric = 0.0



    def train(self, dataloader):
        """
        对一个训练集进行训练，并返回训练集上的平均损失和准确率

        :param dataloader: 训练集数据加载器
        :return: 训练集上的平均损失和准确率
        """
        # 设置模型为训练模式
        self.model.train()
        train_loss = 0
        real_labels = []
        pred_labels = []
        pred_probs = []
        for data, target in tqdm(dataloader):
            # 将数据和标签移动到指定设备上
            data, target = data.to(self.device), target.to(self.device)
            # 将梯度清零
            self.optimizer.zero_grad()
            # 前向传播
            output = self.model(data)

            # 计算损失
            loss = self.criterion(output, target)

            # 反向传播
            loss.backward()
            # 更新权重
            self.optimizer.step()

            if self.scheduler:
                self.scheduler.step()

            # 计算平均损失和准确率
            train_loss += loss.item()
            _pred_labels = torch.argmax(output.logits, dim=-1)
            _pred_probs = torch.softmax(output, dim=-1)
            real_labels += target.cpu().numpy().tolist()
            pred_labels += _pred_labels.cpu().numpy().tolist()
            pred_probs += _pred_probs.cpu().numpy().tolist()

        metric = cal_acc_precision_recall_f1(real_labels, pred_labels, pred_probs=pred_probs, average='binary')
        metric['loss'] = train_loss / len(dataloader)

        return metric

    def evaluate(self, dataloader):
        """
        对一个测试集进行评估，并返回测试集上的平均损失和准确率

        :param dataloader: 测试集数据加载器
        :return: 测试集上的平均损失和准确率
        """
        # 设置模型为评估模式
        self.model.eval()
        valid_loss = 0
        real_labels = []
        pred_labels = []
        pred_probs = []
        # 关闭梯度计算
        with torch.no_grad():
            for data, target in dataloader:
                # 将数据和标签移动到指定设备上
                data, target = data.to(self.device), target.to(self.device)
                # 前向传播
                output = self.model(data)
                # 计算损失
                loss = self.criterion(output, target)

                # 计算平均损失和准确率
                valid_loss += loss.item()
                _pred_labels = torch.argmax(output.logits, dim=-1)
                _pred_probs = torch.softmax(output, dim=-1)
                real_labels += target.cpu().numpy().tolist()
                pred_labels += _pred_labels.cpu().numpy().tolist()
                pred_probs += _pred_probs.cpu().numpy().tolist()

        metric = cal_acc_precision_recall_f1(real_labels, pred_labels, pred_probs=pred_probs, average='binary')
        metric['loss'] = valid_loss / len(dataloader)

        return metric

    def save_best_model(self, save_path, curr_metric):
        model_file_path = os.path.join(save_path, 'weights.pth')
        # 如果当前准确率比最好准确率还要好，就保存当前模型
        if self.best_metric <= curr_metric:
            self.best_metric = curr_metric
            torch.save(self.model.state_dict(), model_file_path)
    
    def train_and_evaluate(self, train_dataloader, valid_dataloader, epochs, model_save_path=None):
        """
        对模型进行训练和评估，并在每个epoch结束时输出训练和测试集上的损失和准确率

        :param train_dataloader: 训练集数据加载器
        :param valid_dataloader: 测试集数据加载器
        :param epochs: 训练轮数
        :param model_save_path: 最佳模型保存的路径（默认为None，表示不保存最佳模型）
        :return: 每个epoch结束时的训练集和测试集上的平均损失和准确率
        """

        for epoch in range(1, epochs + 1):
            train_metric = self.train(train_dataloader)
            valid_metric = self.evaluate(valid_dataloader)


            print(f'Epoch {epoch}: Train metric={train_metric}\nvalid metric={valid_metric}')
            if not os.path.exists(model_save_path):
                os.makedirs(model_save_path)
            self.save_best_model(model_save_path, curr_metric=valid_metric['f1'])

        return None



class PLMTrainer(Trainer):
    def __init__(self, model, optimizer, criterion, scheduler=None, device=None):
        super(PLMTrainer, self).__init__(model, optimizer, criterion, scheduler, device)

    def train(self, dataloader):
        """
        对一个训练集进行训练，并返回训练集上的平均损失和准确率

        :param dataloader: 训练集数据加载器
        :return: 训练集上的平均损失和准确率
        """
        # 设置模型为训练模式
        self.model.train()
        train_loss = 0
        real_labels = []
        pred_labels = []
        pred_probs = []
        for batch in tqdm(dataloader):
            # 将数据和标签移动到指定设备上
            # batch包含了input_ids和attention_mask
            batch = {k: batch[k].to(self.device) for k in batch}
            # 将梯度清零
            self.optimizer.zero_grad()
            # 前向传播
            output = self.model(**batch)
            loss = output.loss

            # 反向传播
            loss.backward()
            # 更新权重
            self.optimizer.step()

            if self.scheduler:
                self.scheduler.step()

            # 计算平均损失和准确率
            train_loss += loss.item()
            # logits是
            _pred_probs = torch.sigmoid(output.logits).squeeze()
            _pred_labels = torch.where(_pred_probs > 0.5, 1, 0)

            real_labels += batch['labels'].cpu().numpy().tolist()
            pred_probs += _pred_probs.detach().cpu().numpy().tolist()
            pred_labels += _pred_labels.cpu().numpy().tolist()
        metric = cal_acc_precision_recall_f1(real_labels, pred_labels, pred_probs=pred_probs, average='binary')

        metric['loss'] = train_loss / len(dataloader)
        return metric

    def evaluate(self, dataloader):
        """
        对一个训练集进行训练，并返回训练集上的平均损失和准确率

        :param dataloader: 训练集数据加载器
        :return: 训练集上的平均损失和准确率
        """
        # 设置模型为训练模式
        # self.model.train()
        self.model.eval()
        valid_loss = 0
        real_labels = []
        pred_labels = []
        pred_probs = []
        for batch in tqdm(dataloader):
            # 将数据和标签移动到指定设备上
            batch = {k: batch[k].to(self.device) for k in batch}

            # 前向传播  
            output = self.model(**batch)
            loss = output.loss


            # 计算平均损失和准确率
            valid_loss += loss.item()
            _pred_probs = torch.sigmoid(output.logits).squeeze()
            _pred_labels = torch.where(_pred_probs > 0.5, 1, 0)

            real_labels += batch['labels'].cpu().numpy().tolist()
            pred_probs += _pred_probs.detach().cpu().numpy().tolist()
            pred_labels += _pred_labels.cpu().numpy().tolist()
        metric = cal_acc_precision_recall_f1(real_labels, pred_labels, pred_probs=pred_probs, average='binary')

        metric['loss'] = valid_loss / len(dataloader)
        return metric

    def save_best_model(self, save_path, curr_metric):
        # 如果当前准确率比最好准确率还要好，就保存当前模型
        if self.best_metric <= curr_metric:
            self.best_metric = curr_metric
            self.model.save_pretrained(save_path)

